import os

from openpyxl import Workbook
from openpyxl.reader.excel import load_workbook


class CeiExcelExportPipeline:
    field_mapping = {
        "title": "标题",
        "time": "时间",
        "source": "来源",
        "content": "内容",
        "url": "地址",
    }

    def __init__(self):
        self.data_buffer = []
        self.batch_size = 5
        self.output_file = 'cei_result.xlsx'

        headers = list(self.field_mapping.keys())
        chinese_headers = [self.field_mapping[h] for h in headers]

        # 检查文件是否存在
        if os.path.exists(self.output_file):
            self.wb = load_workbook(self.output_file)
            self.ws = self.wb.active
        else:
            self.wb = Workbook()
            self.ws = self.wb.active
            self.ws.append(chinese_headers)

    def process_item(self, item, spider):
        if spider.name != "cei":
            return item

        self.add_to_buffer(item)
        if len(self.data_buffer) >= self.batch_size:
            self.flush_buffer()
        return item

    def add_to_buffer(self, item):
        row_data = [item.get(field, "") for field in self.field_mapping.keys()]
        self.data_buffer.append(row_data)

    def flush_buffer(self):
        for row_data in self.data_buffer:
            self.ws.append(row_data)
        self.wb.save(self.output_file)  # 临时保存
        self.data_buffer = []  # 清空缓冲区

    def close_spider(self, spider):
        if spider.name != "cei":
            return
        if self.data_buffer:
            self.flush_buffer()
